package health_system.health_system.websocket;

import health_system.health_system.websocket.Model.ChatMessageHandler;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Pattern;

/**
 * 消息安全工具类
 * 提供消息验证、内容过滤和安全检查功能
 */
public class MessageSecurityUtils {
    private static final Logger logger = Logger.getLogger(MessageSecurityUtils.class.getName());
    
    // 防DoS攻击的速率限制映射（用户ID -> 上次消息时间）
    private static final Map<String, Long> rateLimit = new ConcurrentHashMap<>();
    
    // 用户会话令牌映射（用户ID -> 令牌）
    private static final Map<String, String> sessionTokens = new ConcurrentHashMap<>();
    
    // 内容类型的大小限制（字节）
    private static final Map<String, Long> contentSizeLimits = new HashMap<>();
    
    // 敏感词过滤器
    private static final Pattern sensitiveWordsPattern;
    
    // 常量定义
    private static final String HMAC_ALGORITHM = "HmacSHA256";
    private static final String SECRET_KEY = "pUUWka3xmBDK7DGPj4LCPRvxm9vyYS6B"; // 实际应用中应从配置中读取
    private static final long RATE_LIMIT_INTERVAL = 100; // 消息间隔最小为100毫秒
    private static final int MAX_TEXT_LENGTH = 4000; // 文本最大长度
    
    static {
        // 初始化各类型内容的大小限制
        contentSizeLimits.put("text", 4L * 1024);        // 4KB
        contentSizeLimits.put("image", 10L * 1024 * 1024); // 10MB
        contentSizeLimits.put("video", 50L * 1024 * 1024); // 50MB
        contentSizeLimits.put("audio", 5L * 1024 * 1024);  // 5MB
        contentSizeLimits.put("file", 100L * 1024 * 1024); // 100MB
        
        // 初始化敏感词过滤器（简单示例）
        sensitiveWordsPattern = Pattern.compile(
            // 这里可以添加需要过滤的敏感词的正则表达式
            "(敏感词1|敏感词2|敏感词3)", 
            Pattern.CASE_INSENSITIVE
        );
    }
    
    /**
     * 验证消息内容
     * 
     * @param message 聊天消息对象
     * @return 如果消息有效返回true，否则返回false
     */
    public static boolean validateMessage(ChatMessageHandler.ChatMessage message) {
        if (message == null) {
            return false;
        }
        
        // 验证必要字段
        if (message.getFromUser() == null || message.getType() == null) {
            return false;
        }
        
        // 验证消息类型
        String type = message.getType();
        if (!isValidMessageType(type)) {
            return false;
        }
        
        // 验证消息大小
        if (!validateMessageSize(message)) {
            return false;
        }
        
        // 对于聊天消息，验证接收者
        if (type.equals(ChatMessageHandler.MessageType.CHAT_MESSAGE.getValue()) && message.getToUser() == null) {
            return false;
        }
        
        return true;
    }
    
    /**
     * 检查消息是否为有效的类型
     * 
     * @param type 消息类型
     * @return 如果类型有效返回true，否则返回false
     */
    private static boolean isValidMessageType(String type) {
        return ChatMessageHandler.MessageType.fromString(type) != null;
    }
    
    /**
     * 验证消息大小是否符合限制
     * 
     * @param message 聊天消息
     * @return 如果消息大小符合限制返回true，否则返回false
     */
    public static boolean validateMessageSize(ChatMessageHandler.ChatMessage message) {
        if (message == null || message.getContentType() == null) {
            return false;
        }
        
        String contentType = message.getContentType();
        String content = message.getContent();
        
        // 如果内容为空，直接返回有效
        if (content == null || content.isEmpty()) {
            return true;
        }
        
        // 检查文本内容长度
        if (contentType.equals(ChatMessageHandler.MessageContentType.TEXT.getValue())) {
            return content.length() <= MAX_TEXT_LENGTH;
        }
        
        // 对于Base64编码的内容，估算实际大小
        long estimatedSize = calculateBase64Size(content);
        Long maxSize = contentSizeLimits.get(contentType);
        
        return maxSize != null && estimatedSize <= maxSize;
    }
    
    /**
     * 计算Base64编码字符串的原始大小（字节）
     * 
     * @param base64String Base64编码的字符串
     * @return 估算的原始大小（字节）
     */
    private static long calculateBase64Size(String base64String) {
        if (base64String == null || base64String.isEmpty()) {
            return 0;
        }
        
        // 如果包含Base64头部，移除头部
        int commaIndex = base64String.indexOf(",");
        if (commaIndex > 0 && base64String.startsWith("data:")) {
            base64String = base64String.substring(commaIndex + 1);
        }
        
        // Base64编码的字符串长度大约是原始二进制大小的4/3
        return (long) (base64String.length() * 0.75);
    }
    
    /**
     * 过滤文本中的敏感内容
     * 
     * @param text 原始文本
     * @return 过滤后的文本
     */
    public static String filterSensitiveContent(String text) {
        if (text == null || text.isEmpty()) {
            return text;
        }
        
        // 将敏感词替换为*
        return sensitiveWordsPattern.matcher(text).replaceAll("***");
    }
    
    /**
     * 检查用户消息速率是否超限
     * 
     * @param userId 用户ID
     * @return 如果速率正常返回true，否则返回false
     */
    public static boolean checkRateLimit(String userId) {
        if (userId == null) {
            return false;
        }
        
        long currentTime = System.currentTimeMillis();
        Long lastMessageTime = rateLimit.get(userId);
        
        if (lastMessageTime != null) {
            if (currentTime - lastMessageTime < RATE_LIMIT_INTERVAL) {
                // 消息发送过快
                return false;
            }
        }
        
        // 更新最后消息时间
        rateLimit.put(userId, currentTime);
        return true;
    }
    
    /**
     * 生成用户会话令牌
     * 
     * @param userId 用户ID
     * @return 会话令牌
     */
    public static String generateSessionToken(String userId) {
        if (userId == null) {
            return null;
        }
        
        String token = generateUniqueToken();
        sessionTokens.put(userId, token);
        return token;
    }
    
    /**
     * 验证用户会话令牌
     * 
     * @param userId 用户ID
     * @param token 会话令牌
     * @return 如果令牌有效返回true，否则返回false
     */
    public static boolean validateSessionToken(String userId, String token) {
        if (userId == null || token == null) {
            return false;
        }
        
        String storedToken = sessionTokens.get(userId);
        return storedToken != null && storedToken.equals(token);
    }
    
    /**
     * 生成消息签名
     * 
     * @param message 消息内容
     * @param timestamp 时间戳
     * @return 签名字符串
     */
    public static String generateMessageSignature(String message, long timestamp) {
        try {
            String data = message + timestamp;
            Mac sha256_HMAC = Mac.getInstance(HMAC_ALGORITHM);
            SecretKeySpec secret_key = new SecretKeySpec(SECRET_KEY.getBytes(StandardCharsets.UTF_8), HMAC_ALGORITHM);
            sha256_HMAC.init(secret_key);
            
            byte[] hash = sha256_HMAC.doFinal(data.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(hash);
        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
            logger.log(Level.SEVERE, "生成消息签名失败", e);
            return null;
        }
    }
    
    /**
     * 验证消息签名
     * 
     * @param message 消息内容
     * @param timestamp 时间戳
     * @param signature 待验证的签名
     * @return 如果签名有效返回true，否则返回false
     */
    public static boolean validateMessageSignature(String message, long timestamp, String signature) {
        if (message == null || signature == null) {
            return false;
        }
        
        String calculatedSignature = generateMessageSignature(message, timestamp);
        return calculatedSignature != null && calculatedSignature.equals(signature);
    }
    
    /**
     * 清理过期的速率限制记录
     */
    public static void cleanupRateLimits() {
        long currentTime = System.currentTimeMillis();
        long expiryTime = currentTime - 60000; // 清理一分钟前的记录
        
        rateLimit.entrySet().removeIf(entry -> entry.getValue() < expiryTime);
    }
    
    /**
     * 生成唯一令牌
     */
    private static String generateUniqueToken() {
        return UUID.randomUUID().toString();
    }
    
    /**
     * 检查文件类型安全性
     * 
     * @param mimeType MIME类型
     * @return 如果文件类型安全返回true，否则返回false
     */
    public static boolean isSecureFileType(String mimeType) {
        return FileTransferUtils.isAllowedFileType(mimeType);
    }
} 